import os, json
import pymysql
import oss2
import time
import requests
import logging
import datetime
import jwt
import sys, random
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
import config
from config import DB_CONFIG
from utils import common


class client:
	def __init__(self):
		self.connection = pymysql.connect(host=DB_CONFIG.HOST[0],user=DB_CONFIG.USER[0],password=DB_CONFIG.PASSWORD[0],database=DB_CONFIG.DATABASE)

		

	def get_connection(self):
		# print(self.connection.open, '2'*30)
		try:
			if not self.connection.open:
				self.connection.ping(reconnect=True)
			else:
				self.connection =  pymysql.connect(host=DB_CONFIG.HOST[0],
								user=DB_CONFIG.USER[0],
								password=DB_CONFIG.PASSWORD[0],
								database=DB_CONFIG.DATABASE)
		except:
			self.connection =  pymysql.connect(host=DB_CONFIG.HOST[0],
								user=DB_CONFIG.USER[0],
								password=DB_CONFIG.PASSWORD[0],
								database=DB_CONFIG.DATABASE)

				


	def fetch_data_from_db(self, sql):
		try:
			self.get_connection()
			cur = self.connection.cursor()
			cur.execute(sql)
			result = [item for item in cur.fetchall()]
			
			# self.connection.close()
			return result
		except Exception as e:
			logging.warning(f'error happened when fetch data : {e}')
			self.get_connection()
			cur = self.connection.cursor()
			cur.execute(sql)
			result = [item for item in cur.fetchall()]
			
			# self.connection.close()
			return result

	

	def update_db(self, sql):
		self.get_connection()
		cur = self.connection.cursor()
		cur.execute(sql)
		self.connection.commit()
		cur.close()
		# self.connection.close()
		

	def list_available_role(self, user_id):
		sql = f'SELECT chat_roles_json FROM `chatgpt_users` where id="{user_id}"'
		roles = eval(self.fetch_data_from_db(sql)[0][0])
		role_str = '("' + '","'.join(roles) + '")'
		## 获取角色
		sql = f'SELECT id, name, marker, icon, hello_msg FROM `chatgpt_chat_roles` where marker in {role_str} '

		data = self.fetch_data_from_db(sql)
		rdict = list()
		for sdata in data:
			rdict.append({key:value for key, value in zip(['id', 'name', 'key', 'icon', 'hello_msg'], sdata)})
		return rdict

	def list_all_role(self):
		sql = f'SELECT id, name, marker, icon, hello_msg FROM `chatgpt_chat_roles`'
		data = self.fetch_data_from_db(sql)
		rdict = list()
		for sdata in data:
			rdict.append({key:value for key, value in zip(['id', 'name', 'key', 'icon', 'hello_msg'], sdata)})
		return rdict

	def update_available_role(self, params):
		sql = f"update chatgpt_users set chat_roles_json='{json.dumps(params.get('keys'))}' where id={params.get('id')}"
		logging.info(f'update role attribute :{sql}')
		return self.update_db(sql)	
		
	def list_available_model(self, enable):
		sql = f'SELECT id,name FROM `chatgpt_chat_models` where enabled="{enable}"'

		data = self.fetch_data_from_db(sql)
		rdict = list()
		for sdata in data:
			rdict.append({key:value for key, value in zip(['id', 'name'], sdata)})
		return rdict

	def get_chat_history(self,chat_id):
		sql = f'select type, icon, content, tokens, created_at from  `chatgpt_chat_history` where chat_id="{chat_id}"'
		data = self.fetch_data_from_db(sql)
		rdict = list()
		for sdata in data:
			rdict.append({key:value for key, value in zip(['type', 'icon', 'content', 'tokens', 'created_at'], sdata)})
		return rdict
	
	def list_chat(self, user_id):
		## 根据userid获取chatid，再根据chatid获取chatrole
		sql = f'SELECT role_id, model_id, chat_id, title, chatgpt_chat_roles.icon AS model FROM `chatgpt_chat_items` inner join `chatgpt_chat_roles` on chatgpt_chat_items.role_id=chatgpt_chat_roles.id where user_id="{user_id}" ORDER BY ChatGPT_chat_items.created_at DESC;'
		#sql = f'SELECT chat_id, title, model FROM `chatgpt_chat_items` where user_id="{user_id}"'
		
		data = self.fetch_data_from_db(sql)
		rdict = list()
		for sdata in data:
			rdict.append({key:value for key, value in zip(['role_id', 'model_id','chat_id', 'title',  'icon'], sdata)})
		return rdict
			
	def get_user_from_id(self, id):
		sql = f'SELECT * FROM `chatgpt_users` where id="{id}" '
		return self.fetch_data_from_db(sql)

	def insert_info(self, params):
		
		sql = """INSERT INTO chatgpt_ai_infos 
			(tags, date, name, detail, url, author, contributor) VALUES (%(tags)s, %(date)s, %(name)s, %(detail)s, %(url)s, %(author)s, %(contributor)s)
			"""
		self.get_connection()
		cur = self.connection.cursor()
		cur.execute(sql, params)
		self.connection.commit()
		return True

	def delete_info(self, params):
		sql = f"""delete from chatgpt_ai_infos where id={params.get('id')}"""
		self.get_connection()
		cur = self.connection.cursor()
		cur.execute(sql, params)
		self.connection.commit()
		return True

	def get_info(self, params):
		sql = f"""select * from chatgpt_ai_infos limit {int(params.get('pageSize'))} offset {(int(params.get('pageNum'))-1) * int(params.get('pageSize'))}
			"""
		data = self.fetch_data_from_db(sql)
		rdict = list()
		
		for sdata in data:
			sdata  = list(sdata)
			sdata[1] = json.loads(sdata[1])
			sdata[2] = sdata[2].strftime('%Y-%m-%d')
			rdict.append({key:value for key, value in zip(['id', 'tags','date', 'name',  'detail', 'url', 'author', 'contributor'], sdata)})
		return rdict

	def update_info(self, params):
		sql = f"""update chatgpt_ai_infos set tags='{json.dumps(params.get("tags"))}', date='{params.get("date")}', name='{params.get("name")}', detail='{params.get("detail")}', url='{params.get("url")}', author='{params.get("author")}', contributor='{params.get("contributor")}' where id={params.get("id")}
			"""
		logging.info(f'update info : {sql}')
		self.update_db(sql)
	
	def register_user(self,params):
		## 用户注册
		salt = common.generate_salt()
		## 检查邀请码
		if params.get('code'):
			sql = f"""select code from chatgpt_invite_codes where code='{params.get("code")}'"""
			data = self.fetch_data_from_db(sql)
			if not data:
				return '邀请码无效'
		##创建用户
		new_user_sql = """INSERT INTO chatgpt_users 
			(username, nickname, password, avatar, salt, total_tokens, tokens, calls, img_calls, expired_time, status, chat_config_json, chat_roles_json, chat_models_json, last_login_at, vip, last_login_ip, created_at, updated_at) 
			VALUES (%(username)s, %(nickname)s, %(password)s, %(avatar)s, %(salt)s, %(total_tokens)s, %(calls)s, %(tokens)s, %(img_calls)s, %(expired_time)s, %(status)s, %(chat_config_json)s, %(chat_roles_json)s, %(chat_models_json)s, %(last_login_at)s, %(vip)s, %(last_login_ip)s, %(created_at)s, %(updated_at)s)
			"""
		
		## 构造数据
		params['nickname'] = params.get('nickname')
		params['username'] = params.get('username')
		params['password'] = common.gen_password(params.get('repass'), salt)
		params['avatar'] = random.choice(common.get_headimg())
		params['salt'] = salt
		params['total_tokens'] = 0
		params['tokens'] = 0
		params['calls'] = 10 ## 初始注册时
		params['img_calls'] = 10
		params['expired_time'] = 0
		params['status'] = True
		params['chat_config_json'] = True
		params['chat_roles_json'] = json.dumps(['gpt']) #默认只有通用机器人
		params['chat_models_json'] = True
		params['vip'] = False
		params['last_login_at'] = 0
		params['last_login_ip'] = 0
		params['created_at'] = datetime.datetime.now()
		params['updated_at'] = datetime.datetime.now()

		self.get_connection()
		cur = self.connection.cursor()
		cur.execute(new_user_sql, params)
		self.connection.commit()
		return True

	def get_role_setting(self, params):
		sql = f"""select context_json from chatgpt_chat_roles where id={params.get('role_id')}"""
		logging.info(f'get role profile : {sql}')
		return self.fetch_data_from_db(sql)

	
	def get_history_for_context(self, params, num=10):
		sql = f"""select type, content from chatgpt_chat_history where chat_id='{params.get('chat_id')}' order by created_at desc limit {num}"""
		logging.info(f'get history data : {sql}')
		return self.fetch_data_from_db(sql)

	def get_apikey_for_model(self, params):
		sql = f"""select value from chatgpt_api_keys where name='{params.get("model")}'"""
		return self.fetch_data_from_db(sql)



	def login(self, username, password):
		sql = f'SELECT salt, password, id, avatar, chat_roles_json, username, nickname, calls, tokens, img_calls, total_tokens, expired_time FROM `chatgpt_users` where username = "{username}" '
		result = self.fetch_data_from_db(sql)
		if not result:
			return False
		else:
			password = common.gen_password(password, result[0][0])
			if password == result[0][1]:
				## TODO 记录日志，生成token
				payload = {
					'sub': username,  # 主题（通常是用户ID或用户名）
					'iat': datetime.datetime.utcnow(),  # 签发时间
					'exp': datetime.datetime.utcnow() + datetime.timedelta(hours=1),  # 过期时间
				}
				# 生成JWT
				token = jwt.encode(payload, config.SECRET_KEY, algorithm='HS256')
				return {"id":result[0][2], "token":token , "avatar":result[0][3], 'chat_roles':json.loads(result[0][4]), 'username':result[0][5], 'nickname':result[0][6], 'calls': result[0][7], 'tokens':result[0][8], 'img_calls':result[0][9], 'total_tokens':result[0][10], 'expired_time':result[0][11]}
			else:
				#密码不对
				return False


	
	def get_config(self, key):
		sql = f'select config_json from `chatgpt_configs` where marker="{key}"'
		return self.fetch_data_from_db(sql)[0][0]


	def update_sd_progressid(self, task_id, progress_id):
		sql = f'update chatgpt_sd_jobs set progress_id="{progress_id}" where task_id="{task_id}"'
		self.update_db(sql)

	def update_image(self,task_id, img_url):
		sql = f'update chatgpt_sd_jobs set img_url="{img_url}" where task_id="{task_id}"'
		self.update_db(sql)
	

	def get_progress(self, task_id):
		sql = f'select progress_id from chatgpt_sd_jobs where task_id="{task_id}"'
		progress_id = self.fetch_data_from_db(sql)[0][0]
		return progress_id
		

	def get_version(self, model):
		sql = f'select value from chatgpt_sd_models where name="{model}"'
		version = self.fetch_data_from_db(sql)[0][0]
		return version


	def download_to_local(self, image_url):
		os.makedirs('cache_imgs', exist_ok=True)
		res = requests.get(image_url)
		if res.status_code == 200:
			wname = f"cache_imgs/{image_url.split('/')[-1]}"
			with open(wname, 'wb') as w:
				w.write(res.content) 
			return wname
		else:
			logging.warning('文件下载失败')


	def uploadfile(self, image_url):
		# lfile = 'cache_imgs/out-0.png'
		lfile = self.download_to_local(image_url)
		# 阿里云账号AccessKey拥有所有API的访问权限，风险很高。强烈建议您创建并使用RAM用户进行API访问或日常运维，请登录RAM控制台创建RAM用户。
		auth = oss2.Auth(config.OSS_KEY,config.OSS_SECRET)
		with open(lfile, 'rb') as fileobj:
			bucket = oss2.Bucket(
				auth, config.OSS_ENDPOINT, config.OSS_BUCKET )
			timestamp = time.strftime('%y%m%d%H%M%S',time.localtime())
			bucket.put_object(f'sd/{timestamp}.png', fileobj)
			return f'{config.OSS_URL}/sd/{timestamp}.png'

	
	def save_chat_history(self, params):
		keys = 'user_id,chat_id,type,icon,role_id,model,content,tokens,use_context,created_at,updated_at,deleted_at'
		params = {key:params.get(key) for key in keys.split(',')}
		# save_sql = "INSERT INTO chatgpt_chat_history (user_id, chat_id, type, icon, role_id, model, content, tokens, use_context, created_at, updated_at, deleted_at) VALUES (%s,'%s','%s','%s',%s,'%s','%s',%s,%s,'%s','%s','%s') "
		# print(params)
		# sql = save_sql % (params.get('user_id', 'NULL'), params.get('chat_id', 'NULL'),params.get('type', 'NULL'),params.get('icon', 'NULL'),params.get('role_id',1),params.get('model', 'NULL'),params.get('content', 'NULL'),params.get('tokens',0),params.get('use_context', 1),params.get('created_at', 'NULL'),params.get('update_at', 'NULL'),params.get('deleted_at', 'NULL'))
		sql = """INSERT INTO ChatGPT_chat_history 
			(user_id, chat_id, type, icon, role_id, model, content, tokens, use_context, created_at, updated_at, deleted_at) 
			VALUES (%(user_id)s, %(chat_id)s, %(type)s, %(icon)s, %(role_id)s, %(model)s, %(content)s, %(tokens)s, %(use_context)s, %(created_at)s, %(updated_at)s, %(deleted_at)s)
			"""
		print(params)
		logging.info(f'存储对话历史:{sql}, 数据:{params}')
		cur = self.connection.cursor()
		cur.execute(sql, params)
		self.connection.commit()

	def get_role_icon_by_id(self, id):
		sql = f"select icon from chatgpt_chat_roles where id='{id}'"
		return self.fetch_data_from_db(sql)

	## 储存对话item，表现为列表中的新一项
	def save_chat_item(self, params):

		#检查chatid是否存在，
		check_sql = f"""select id from chatgpt_chat_items where chat_id='{params.get("chat_id")}' """
		logging.info(f'add chat item or not : {check_sql}')
		data = self.fetch_data_from_db(check_sql)
		if data:
			logging.info('对话item已存在，不再存储')
			return 


		keys = 'user_id,chat_id,role_id,model_id,model,title,created_at,updated_at,deleted_at'
		params = {key:params.get(key) for key in keys.split(',')}

		sql = """INSERT INTO chatgpt_chat_items 
			(user_id,chat_id,role_id, model_id, model, title, created_at,updated_at,deleted_at) 
			VALUES (%(user_id)s, %(chat_id)s, %(role_id)s, %(model_id)s, %(model)s, %(title)s, %(created_at)s, %(updated_at)s, %(deleted_at)s)
			"""
		logging.info(f'存储对话item:{sql}, 数据:{params}')
		cur = self.connection.cursor()
		cur.execute(sql, params)
		self.connection.commit()

	def check_calls(self, params):
		sql = f"select calls from chatgpt_users where id={params.get('user_id')}"
		logging.info(f' 获取后台用户调用次数 : {sql}')
		calls = self.fetch_data_from_db(sql)
		if not calls:
			## 非注册用户
			sql = f"select chat_id from chatgpt_chat_history where user_id={params.get('user_id')} and type='reply' "
			logging.info(f' 获取后台用户已调用次数 : {sql}')
			return 1 - len(sql)## 只允许一次
		else:
			return calls[0][0]


	def reset_calls(self, params):
		## 获取模型权重
		sql = f'''select weight from chatgpt_chat_models where value='{params.get("model")}' '''
		logging.info(f'get model weight : {sql}')
		weight = self.fetch_data_from_db(sql)[0][0]
		## 更新
		print(params)
		update_sql = f"update chatgpt_users set calls=calls-{weight} where id={params.get('model_id')} "
		logging.info(f'update calls : {update_sql}')
		self.update_db(update_sql)

	def reset_sd_calls(self, params):
		## 获取模型权重
		sql = f'''select weight from chatgpt_sd_models where value='{params.get("model")}' '''
		logging.info(f'get model weight : {sql}')
		weight = self.fetch_data_from_db(sql)[0][0]
		## 更新
		update_sql = f"update chatgpt_users set img_calls=img_calls-{weight} where id={params.get('model_id')} "
		logging.info(f'update calls : {update_sql}')
		self.update_db(update_sql)


	def get_tokens(self, params):
		## 获取模型消耗token
		sql = f"""select * from chatgpt_chat_history where chat_id='{params.get("chat_id")}' and type='reply' order by created_at desc limit 1"""
		logging.info(f'获取回复的总token:{sql}')
		return self.fetch_data_from_db(sql)

	
	def get_model_by_id(self, mid):
		sql = f'select name from `chatgpt_chat_models` where id={mid}'
		data = self.fetch_data_from_db(sql)
		return data[0][0]

	def update_chat_title(self, params):
		sql = f"update chatgpt_chat_items set title='{params.get('title')}' where chat_id='{params.get('chat_id')}'"
		self.update_db(sql)
		

	def remove_chat_by_id(self, params):
		sql = f"delete from chatgpt_chat_items where chat_id='{params.get('chat_id')}'"
		self.update_db(sql)

	def clear_chat_by_id(self, params):
		sql = f"delete from chatgpt_chat_items where user_id='{params.get('user_id')}'"
		self.update_db(sql)
		
		
	def get_mj_jobs(self, params):
		self.get_connection()
		cur = self.connection.cursor()

		if params.get('status') == '1':
			sql = 'select * from chatgpt_mj_jobs where progress=100'
		else:
			sql = 'select * from chatgpt_mj_jobs where progress<100'

		if params.get('user_id'):
			sql += f" and user_id={params.get('user_id')} " 

		if params.get('publish'):
			sql += f" and publish={params.get('publish')} " 

		if params.get('page') and params.get('page_size'):
			offset = (int(params.get('page',1)) - 1) * int(params.get('page_size',0))
			page_size = int(params.get('page_size',0))
			sql += f' limit {page_size} offset {offset} '


		logging.info(f'查询jobs : {sql}')
		result = self.fetch_data_from_db(sql)
		return result

	def get_sd_jobs(self, params):
		self.get_connection()
		cur = self.connection.cursor()

		if params.get('status') == '1':
			sql = 'select * from chatgpt_sd_jobs  where progress=100'
		else:
			sql = 'select * from chatgpt_sd_jobs where progress<100'

		if params.get('user_id'):
			sql += f" and user_id={params.get('user_id')} " 

		if params.get('publish'):
			sql += f" and publish={params.get('publish')} " 

		if params.get('page') and params.get('page_size'):
			offset = (int(params.get('page',1)) - 1) * int(params.get('page_size',0))
			page_size = int(params.get('page_size',0))
			sql += f' limit {page_size} offset {offset} '


		logging.info(f'查询jobs : {sql}')
		result = self.fetch_data_from_db(sql)
		return result

	def get_sd_models(self, params):
		self.get_connection()
		cur = self.connection.cursor()
		
		sql = 'select name from chatgpt_sd_models where enabled=1'
		
		logging.info(f'查询jobs : {sql}')
		result = self.fetch_data_from_db(sql)
		return result

	def create_sd_job(self, data):
		self.get_connection()
		cur = self.connection.cursor()
		sql = """INSERT INTO chatgpt_sd_jobs 
			(user_id,type,task_id,prompt,img_url,params,progress,publish,err_msg,created_at,progress_id) 
			VALUES (%(user_id)s, %(type)s, %(task_id)s, %(prompt)s, %(img_url)s, %(params)s, %(progress)s, %(publish)s, %(err_msg)s, %(created_at)s, %(progress_id)s)
			"""
		logging.info(f'生成绘图任务:{sql}, 数据:{data}')
		cur = self.connection.cursor()
		cur.execute(sql, data)
		self.connection.commit()


	def update_sd_job_progress(self, task_id, progress):
		update_sql = f"update chatgpt_sd_jobs set progress={progress*100} where task_id='{task_id}'"
		logging.info(f'update progress {update_sql}')
		self.update_db(update_sql)
		


	def remove_sd_job(self, data):
		sql = f"delete from chatgpt_sd_jobs where id={data.get('id')}"
		self.update_db(sql)
		
	
	def publish_sd_job(self, data):
		sql = f"update chatgpt_sd_jobs set publish={data.get('action')} where id={data.get('id')}"
		logging.info(f'删除绘图任务:{sql}')
		self.update_db(sql)

	
	def get_order_list(self, params):

		offset = (int(params.get('page',1)) - 1) * int(params.get('page_size',0))
		page_size = int(params.get('page_size',0))
		sql = f"select order_no,remark,pay_time from chatgpt_orders where user_id={params.get('id')} and status=2 order by id DESC limit {page_size} offset {offset} "
		logging.info(f'获取订单列表:{sql}')
		return self.fetch_data_from_db(sql)

	def get_product_list(self, params=None):
		sql = f"select name,price,discount,days,calls,img_calls,enabled,sales,sort_num from chatgpt_products where enabled=1 order by sort_num ASC"
		logging.info(f'获取产品列表:{sql}')
		return self.fetch_data_from_db(sql)


	def get_invite_code(self, params):

		sql = f"select hits, reg_num, code FROM `chatgpt_invite_codes` where user_id={params.get('id')}"
		data = self.fetch_data_from_db(sql)
		if not data:
			create_code_sql = f"insert INTO chatgpt_invite_codes (user_id, code, hits, reg_num, created_at) VALUES (%(user_id)s, %(code)s, %(hits)s, %(reg_num)s, %(created_at)s)"

			params['user_id'] = params['id']
			params.pop('id')
			params['code'] = common.generate_salt(8)
			params['hits'] = 0
			params['reg_num'] = 0
			params['created_at'] = datetime.datetime.now()

			data = [[0,0,params['code']]]
			cur = self.connection.cursor()
			cur.execute(create_code_sql, params)
			self.connection.commit()

		return data


		
	def get_invite_list(self, params):
		offset = (int(params.get('page',1)) - 1) * int(params.get('page_size',0))
		page_size = int(params.get('page_size',0))
		sql = f"select reward_json, created_at from chatgpt_invite_logs where user_id={params.get('id')} limit {page_size} offset {offset} "
		logging.info(f'获取邀请列表:{sql}')
		return self.fetch_data_from_db(sql)

		



if __name__ == '__main__':	
	db = client()
	db.get_model_by_id(1)